// Lasso-LDA
// (c) Yongjin Park, 2013

#include <sstream>
#include <boost/lexical_cast.hpp>
#include "util.hh"
#include "lda.hh"
#include "lasso.hh"
#include "expression.hh"
#include <string>
#include <iostream>
#include <fstream>

typedef boost::ptr_vector<func_lasso_pred_t> lasso_pred_vec_t;
bool read_lasso_pred_vec(const char* beta_file, const char* prec_file,
    const char* hyper_file, lasso_pred_vec_t& out)
{
  size_t n = num_rows(beta_file);
  size_t p = num_columns(beta_file);

  if (n != num_rows(prec_file))
    return false;

  if (p * p != num_columns(prec_file))
    return false;

  if (n != num_rows(hyper_file))
    return false;

  typedef func_lasso_pred_t::vec_type vec_t;
  typedef func_lasso_pred_t::mat_type mat_t;

  std::string betaline;
  std::string precline;
  std::string hyperline;

  mat_t prec;
  vec_t beta(p, 0.);
  for (size_t j = 0; j < p; ++j)
    prec.push_back(new vec_t(p, 0.));

  out.clear();

  std::ifstream beta_in(beta_file, std::ios::in);
  std::ifstream prec_in(prec_file, std::ios::in);
  std::ifstream hyper_in(hyper_file, std::ios::in);

  for (size_t i = 0;
      (i < n) && std::getline(beta_in, betaline)
          && std::getline(prec_in, precline)
          && std::getline(hyper_in, hyperline); ++i)
  {
    // read each line & fill in beta and covariance
    std::istringstream beta_iss(betaline);
    std::istringstream prec_iss(precline);
    std::istringstream hyper_iss(hyperline);

    double penalty, a0, b0;
    hyper_iss >> penalty >> a0 >> b0;

    double x;
    for (size_t j = 0; (j < p) && (beta_iss >> x); ++j)
      beta[j] = x;

    for (size_t j = 0; j < p; ++j)
      for (size_t k = 0; (k < p) && (prec_iss >> x); ++k)
        prec[j][k] = x;

    out.push_back(new func_lasso_pred_t(a0, b0, beta, prec));

//    for (size_t j = 0; j < p; ++j)
//      dump_vec(prec[j]);
  }

  beta_in.close();
  prec_in.close();
  hyper_in.close();

  return true;
}


struct prog_args_t
{

  explicit prog_args_t() :
      g_file(""), x_file(""), y_file(""), // required args
      mean_file(""), prec_file(""), hyper_file(""), //
      output("output"), sample(100)       // options
  {
  }

  std::string g_file;
  std::string x_file;
  std::string y_file;
  std::string mean_file;
  std::string prec_file;
  std::string hyper_file;
  std::string output;
  size_t sample;
};

void print_help(const char* prog, const prog_args_t& args)
{
  std::cerr << prog << " -g G.txt -x X.txt -y Y.txt -o out" << std::endl;
  std::cerr << "         -mean beta.mean -prec beta.prec -hyper beta.hyper"
      << std::endl;
  std::cerr << "         -s sample" << std::endl;
  std::cerr << std::endl;
  std::cerr
      << " G.txt  geneset / group file (each line contains row-indexes of X,Y)"
      << std::endl;
  std::cerr << " X.txt      control expression file (n x p tab-sep real mat)"
      << std::endl;
  std::cerr << " Y.txt      case expression file (n x p tab-sep real mat)"
      << std::endl;
  std::cerr << " beta.mean  mean coefficient file (result of fitting)"
      << std::endl;
  std::cerr << " beta.prec  precision of coefficient file (result of fitting)"
      << std::endl;
  std::cerr << " beta.hyper hyper-paramter file (result of fitting)"
      << std::endl;
  std::cerr << " out        output files' header (default: " << args.output
      << ")" << std::endl;
  std::cerr << " sample     number of samples (default: " << args.sample << ")"
      << std::endl;
  std::cerr << std::endl;
}

template<typename T>
T get_param(const char* arg, std::string emsg)
{
  T val;
  try
  {
    val = boost::lexical_cast<T>(arg);
  } catch (boost::bad_lexical_cast&)
  {
    err_msg(emsg);
    exit(1);
  }
  return val;
}

bool parse_args(const int argc, const char* argv[], prog_args_t& args_out)
{
  for (int j = 1; j < argc; ++j)
  {
    std::string curr = argv[j];

    if (curr == "-g" && ++j < argc)
    {
      args_out.g_file = argv[j];
    }
    else if (curr == "-x" && ++j < argc)
    {
      args_out.x_file = argv[j];
    }
    else if (curr == "-y" && ++j < argc)
    {
      args_out.y_file = argv[j];
    }
    else if (curr == "-o" && ++j < argc)
    {
      args_out.output = argv[j];
    }
    else if (curr == "-mean" && ++j < argc)
    {
      args_out.mean_file = argv[j];
    }
    else if (curr == "-prec" && ++j < argc)
    {
      args_out.prec_file = argv[j];
    }
    else if (curr == "-hyper" && ++j < argc)
    {
      args_out.hyper_file = argv[j];
    }
    else if (curr == "-s" && ++j < argc)
    {
      args_out.sample = get_param<size_t>(argv[j], "parsing #samples");
    }

  }

  if (args_out.g_file.size() == 0 || args_out.x_file.size() == 0
      || args_out.y_file.size() == 0 || args_out.mean_file.size() == 0
      || args_out.prec_file.size() == 0 || args_out.hyper_file.size() == 0)
    return false;

  return true;
}

int main(const int argc, const char* argv[])
{
  prog_args_t args;
  if (!parse_args(argc, argv, args))
  {
    print_help(argv[0], args);
    return 1;
  }

  expr_lm_vec_t data_vec;
  size_t n, p;
  const char* x_file = args.x_file.c_str();
  const char* y_file = args.y_file.c_str();

  if (!read_expr_lm_vec(x_file, y_file, data_vec, n, p))
    return -1;

  const char* beta_file = args.mean_file.c_str();
  const char* prec_file = args.prec_file.c_str();
  const char* hyper_file = args.hyper_file.c_str();

  lasso_pred_vec_t topics;
  read_lasso_pred_vec(beta_file, prec_file, hyper_file, topics);
  const size_t K = topics.size();

  TLOG("read " << K << " topics");

  // read geneset data
  std::string& g = args.g_file;
  boost::ptr_vector<geneset_t<expr_lm_t> > genesets;
  double a0 = 1. / ((double) K);
  if (!read_genesets(g.c_str(), data_vec, a0, K, genesets))
    return -1;

  TLOG("read " << genesets.size() << " genesets");

  gibbs_latent_update_t<expr_lm_t, func_lasso_pred_t> gibbs_update(topics, args.sample);
  std::for_each( genesets.begin(), genesets.end(), gibbs_update );

  output_argmax(args.output, data_vec, topics);
  output_genesets(args.output, genesets);
  std::ofstream ofs((args.output + ".llik").c_str(), std::ios::out);
  ofs << likelihood(genesets) << std::endl;
  ofs.close();

  return 0;
}
